from transformers import BertTokenizer, BertModel
import torch
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

# 加载 tokenizer 和模型
model_name = 'D:\\work\program\\pytorch_models\\bert-base-chinese'
tokenizer = BertTokenizer.from_pretrained(model_name)
model = BertModel.from_pretrained(model_name)

# 输入长文本
text = """
近年来，随着人工智能技术的快速发展，自然语言处理（NLP）领域取得了显著的进展。特别是深度学习模型的应用，使得机器能够更好地理解和生成自然语言。其中，BERT（Bidirectional Encoder Representations from Transformers）是一个典型的例子，它通过预训练大量文本数据，学会了捕捉文本中的深层语义信息。此外，还有一些专门用于文本摘要的模型，如BART和T5，这些模型不仅能够生成高质量的摘要，还能应用于多种自然语言处理任务。
"""

text2 = """
客户:我想找个快递 /n 助手:请您提供一下快递单号或单号后四位，我帮您查询一下 /n 客户:6665 /n 助手:我会通知快递小哥核实情况，让他尽快给您回复 /n 客户:我有个快递放驿站了，但是我签收到取货码 /n 助手:请提供一下快递单号后四位 /n 客户:6665 /n 助手:未查询到该单号的取货码，请核实一下单号是否正确 /n 客户:好的，再见 /n 助手:好的，再见
"""

# 分句
sentences = text2.split('。')

# 编码句子
encoded_sentences = [tokenizer.encode(sentence, add_special_tokens=True, max_length=512, truncation=True) for sentence in sentences]

# 获取句向量
with torch.no_grad():
    vectors = []
    for encoded_sentence in encoded_sentences:
        input_ids = torch.tensor([encoded_sentence])
        outputs = model(input_ids)
        vector = outputs.last_hidden_state.mean(dim=1).squeeze()
        vectors.append(vector.numpy())

# 计算句子的重要性
tfidf_vectorizer = TfidfVectorizer()
tfidf_matrix = tfidf_vectorizer.fit_transform(sentences)


def summaries1():
    # 直接在稀疏矩阵上计算每行的和
    sentence_importance = tfidf_matrix.sum(axis=1).A1  # .A1 将结果转换为一维 NumPy 数组
    # 选择最重要的句子
    top_n = 3  # 选择前3个最重要的句子
    top_indices = sentence_importance.argsort()[-top_n:][::-1]
    summary_sentences = [sentences[i] for i in top_indices]
    # 生成摘要
    summary = '。'.join(summary_sentences) + '。'
    print(summary)

def summaries2():
    feature_names = tfidf_vectorizer.get_feature_names_out()
    # 获取每个文档的最高TF-IDF分数的词汇
    top_n = 3  # 您想要包含在摘要中的最高TF-IDF分数的词汇数量
    summaries = []
    for doc_idx in range(tfidf_matrix.shape[0]):
        tfidf_scores = tfidf_matrix[doc_idx, :].toarray().flatten()
        top_indices = np.argsort(tfidf_scores)[-top_n:][::-1]  # 获取分数最高的索引
        summary = [feature_names[i] for i in top_indices]
        summaries.append(" ".join(summary))
    # 输出摘要
    for i, summary in enumerate(summaries):
        print(f"Document {i + 1} summary: {summary}")

# summaries1()

summaries2()